import pandas as pd
import re
import numpy as np
import sys


def gen_word_list(raw_list):
    word_list = []
    for sentence in raw_list:
        sentence = re.sub(r"\W+", " ", sentence).lower()
        word_list += sentence.split()
    return word_list


def gen_word_count(word_list):
    word_count_dict = {}
    for word in word_list:
        try:
            word_count_dict[word] += 1
        except KeyError:
            word_count_dict[word] = 1
    return word_count_dict


def gen_probs(word_dict):
    probs = {x: np.log(word_dict[x] / sum(word_dict.values()))
             for x in word_dict}
    return probs


def delete_letter(word):
    """
    Output: a list of all possible strings after one delete
    """
    split_l = [(word[:i], word[i:]) for i in range(len(word) + 1)]
    delete_l = [L + R[1:] for L, R in split_l if R]
    return delete_l


def switch_letter(word):
    """
    output a list of all possible strings after one switched
    """
    split_l = [(word[:i], word[i:]) for i in range(len(word))]
    switch_l = [L + R[1] + R[0] + R[2:] for L, R in split_l if R and len(R) >= 2]
    return switch_l


def replace_letter(word):
    """
    Output a list of all possible strings after one replace.
    """
    letters = r'abcdefghijklmnopqrstuvwxyz'
    split_l = [(word[:i], word[i:]) for i in range(len(word) + 1)]
    replace_set = set()
    for i in letters:
        for l, r in split_l:
            if r and r[0] != i:
                new = l + i + r[1:]
                replace_set.add(new)
    return sorted(list(replace_set))


def insert_letter(word):
    """
    Output a set of all possible strings after one insert
    """
    letters = r'abcdefghijklmnopqrstuvwxyz'
    insert_l = []
    split_l = [(word[:i], word[i:]) for i in range(len(word) + 1)]
    for i in letters:
        for l, r in split_l:
            new = l + i + r
            insert_l.append(new)
    return insert_l


def edit_one_letter(word):
    """
    Output a set of words after one edit.
    """
    edit_one_set = set(insert_letter(word) + delete_letter(word)
                       + replace_letter(word) + switch_letter(word))
    return edit_one_set


def edit_two_letters(word):
    """
    Output a set of strings after two edit
    """
    edit_two_set = set()
    for i in edit_one_letter(word):
        edit_two_set = edit_two_set | edit_one_letter(i)
    return edit_two_set


def auto_correct(word, probs, n=2):
    """
    n: number of correction to return
    """
    if word in probs:
        return [word]
    one = edit_one_letter(word)
    two = edit_two_letters(word)
    sug = one or two
    b = [(x, probs[x]) for x in sug if x in probs]
    suggestions = sorted(b, key=lambda x: float(x[1]), reverse=True)
    n_best = [x[0] for x in suggestions[:n]]
    return n_best


def save_probs(path='data/probs.csv'):
    df_probs = pd.DataFrame(probs, index=[0])
    df_probs.to_csv(path)


def gen_dataset(path='data/df_imda8k_train.csv'):
    df_raw = pd.read_csv(path, index_col=0)
    raw_list = df_raw['raw_text'].to_list()
    word_list = gen_word_list(raw_list)
    word_dict = gen_word_count(word_list)
    probs = gen_probs(word_dict)
    return probs


if __name__ == '__main__':
    try:
        probs = pd.read_csv('data/probs.csv', index_col=0)
        print(probs)
    except FileNotFoundError:
        probs = gen_dataset()
    while 1:
        print('enter a word: ', end=" ")
        word = re.sub(r"\W+", "", sys.stdin.readline())
        if word == '':
            break
        word_list = auto_correct(word, probs, n=2)
        print("Auto Correct Result: ", end=" ")
        for w in word_list:
            print(w, end=" ")
        print()